// you should define SIZEX/SIZEY/SIZEZ/BLKX/BLKY/ITER
// actually 2d-partitioning in this program

// #include <iostream>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <assert.h>

#include <sys/time.h>
#include "aceMesh_runtime_c.h"

#define MULTIG 4
double A[MULTIG][2][SIZEX][SIZEY][SIZEZ]; //__attribute__ ((aligned(64)));
// for test
double test[MULTIG][2][SIZEX][SIZEY][SIZEZ];
int BLKX, BLKY;
#define DIST (1)
// arg
#define alpha_d 0.0876
#define beta_d 0.0765
// kernel
#define Do3d7pkernel(h, dest, src, i, j, k)                                         \
    {                                                                            \
        A[h][dest][i][j][k] = alpha_d * (A[h][(src)][(i)][(j)][(k)]) +                 \
                           beta_d * (A[h][src][i - 1][j][k] + A[h][src][i + 1][j][k] + \
                                     A[h][src][i][j - 1][k] + A[h][src][i][j + 1][k] + \
                                     A[h][src][i][j][k - 1] + A[h][src][i][j][k + 1]); \
    }
//
#define TestDo3d7pkernel(h, itr, i, j, k)                                                                           \
    {                                                                                                            \
        test[h][itr % 2][i][j][k] = alpha_d * test[h][(itr + 1) % 2][i][j][k] +                                        \
                                 beta_d * (test[h][(itr + 1) % 2][i - 1][j][k] + test[h][(itr + 1) % 2][i + 1][j][k] + \
                                           test[h][(itr + 1) % 2][i][j - 1][k] + test[h][(itr + 1) % 2][i][j + 1][k] + \
                                           test[h][(itr + 1) % 2][i][j][k - 1] + test[h][(itr + 1) % 2][i][j][k + 1]); \
    }

typedef struct
{
    int starti, endi;
    int startj, endj;
    int iter;
    int h;
    void *src_addr;
} targs;

typedef struct
{
    int iter;
    int bsx;
    int bsy;
    int h;
    void *src_addr;
} targs_outter;

void my_neighbors(struct ptr_array *neighbor_addrs, targs *args)
{
    int starti = args->starti;
    int startj = args->startj;
    void *src_addr = args->src_addr;

    neighbor_addrs->len = 0;
    if (starti - BLKX > 0)
        neighbor_addrs->arr[(neighbor_addrs->len)++]= (void *)(src_addr - BLKX * SIZEY * SIZEZ * 8);
    if (starti + BLKX < SIZEX-1)
        neighbor_addrs->arr[(neighbor_addrs->len)++]= (void *)(src_addr + BLKX * SIZEY * SIZEZ * 8);
    if (startj - BLKY > 0)
        neighbor_addrs->arr[(neighbor_addrs->len)++]= (void *)(src_addr - BLKY * SIZEZ * 8);
    if (startj + BLKY < SIZEY-1)
        neighbor_addrs->arr[(neighbor_addrs->len)++]= (void *)(src_addr + BLKY * SIZEZ * 8);
    //printf("neighbor_len=%d\n",neighbor_addrs->len);
}
#ifdef INIT_CYCLIC
void stencil_init_cyclic(targs *args)
{
    int starti = args->starti;
    int endi = args->endi;
    int startj = args->startj;
    int endj = args->endj;
    int i, j, k;

    for (i = starti; i < endi; i++)
        for (j = startj; j < endj; j++)
            for (k = 0; k < SIZEZ; k++)
            {
                A[0][i][j][k] = (double)(i * 2.5 + j * 3.3 + k * 0.5) / 3.0;
                A[1][i][j][k] = A[0][i][j][k];
            }
}
#endif
#ifdef INIT_BLOCKED
void stencil_init_blocked(targs *args)
{
    int starti = args->starti;
    int endi = args->endi;
    int i, j, k;

    for (i = starti; i < endi; i++)
        for (j = 0; j < SIZEY; j++)
            for (k = 0; k < SIZEZ; k++)
            {
                A[0][i][j][k] = (double)(i * 2.5 + j * 3.3 + k * 0.5) / 3.0;
                A[1][i][j][k] = A[0][i][j][k];
            }
}
#endif
void stencil_core(targs *args)
{
    int starti = args->starti;
    int endi = args->endi;
    int startj = args->startj;
    int endj = args->endj;
    int iter = args->iter;
    int h = args->h;
    int i, j, k;

    int dest = iter % 2;
    int src = 1 - dest;
    for (i = starti; i < endi; i++)
        for (j = startj; j < endj; j++)
            // for(k = 1; k < args->k_N - 1; k++)
            for (k = 1; k < SIZEZ - 1; k++)
                Do3d7pkernel(h, dest, src, i, j, k);
}

int reg_3d7p(targs_outter *outer_each)
{
    targs each;
    int itr, i, j;

    int blocsizex = outer_each->bsx;
    int blocsizey = outer_each->bsy;
    int outer_iter = outer_each->iter;
    int h=outer_each->h;

    int index=0;
    for (int inner_iter = 0; inner_iter < DIST; inner_iter++)
    {
        int local_iter = outer_iter + inner_iter;
        acemesh_begin_split_task("3d7p iter");

        for (i = 1; i < SIZEX - 1; i += blocsizex)
        {
            int endi = (i + blocsizex > SIZEX - 1) ? SIZEX - 1 : i + blocsizex;
            for (j = 1; j < SIZEY - 1; j += blocsizey)
            {
                //printf("register %d subtask begin\n", index);
                //fflush(stdout);
                //                int indexj = (j-1) / blocsizey;
                int endj = (j + blocsizey > SIZEY - 1) ? SIZEY - 1 : j + blocsizey;
                each.starti = i;
                each.endi = endi;
                each.startj = j;
                each.endj = endj;
                each.iter = local_iter;
                each.h = h;
                each.src_addr = &A[h][(local_iter + 1) % 2][i][j][0];
                acemesh_push_wlist(1, &A[h][local_iter % 2][i][j][0], NORMAL);
                acemesh_push_rlist(1, &A[h][(local_iter + 1) % 2][i][j][0], NORMAL);
                acemesh_task_generator_with_neighbors((TASK_FUNCPTR)stencil_core,
                                                      (void *)(&each), sizeof(targs),
                                                      NULL, (NEIGHBOR_FUNCPTR)my_neighbors, &each);
                acemesh_task_set_type(STENCIL_TASK);
                //printf("register %d subtask end\n", index++ );
                //fflush(stdout);
            }
        }
        acemesh_end_split_task();
    }
    acemesh_spawn_and_goout();
    return 0;
}

int Do3d7p(const int iter, const int blocsizex, const int blocsizey)
{
    targs_outter each;
    int itr, inner_itr, i, j, h;

    for(h=0; h<MULTIG; h++)
    {
        for (itr = 0; itr < iter; itr += DIST)
        {
            
            acemesh_begin_split_task("3d7p iter");

            for (int outter_iter = 0; outter_iter < DIST; outter_iter++) //当DIST>2时，存在重复注册的地址，会有环？
            {
                inner_itr=itr+outter_iter;
                for (i = 1; i < SIZEX - 1; i += blocsizex)
                {
                    int endi = (i + blocsizex > SIZEX - 1) ? SIZEX - 1 : i + blocsizex;
                    for (j = 1; j < SIZEY - 1; j += blocsizey)
                    {
                        int endj = (j + blocsizey > SIZEY - 1) ? SIZEY - 1 : j + blocsizey;
                        acemesh_push_wlist(1, &A[h][inner_itr % 2][i][j][0], NORMAL);
                        acemesh_push_rlist(1, &A[h][(inner_itr + 1) % 2][i][j][0], NORMAL);
                    //acemesh_push_weak_wlist(1, &A[inner_itr % 2][i][j][0], NORMAL);
                    //acemesh_push_weak_rlist(1, &A[(inner_itr + 1) % 2][i][j][0], NORMAL);
                    }
                }
            }
            each.iter = itr;
            each.h = h;
            each.bsx = blocsizex;
            each.bsy = blocsizey;
            //acemesh_composite_task_generator((TASK_FUNCPTR)reg_3d7p, (void *)(&each), sizeof(targs_outter));
            acemesh_task_generator((TASK_FUNCPTR)reg_3d7p, (void *)(&each), sizeof(targs_outter));
            //acemesh_task_set_type(COMPOSITE_TASK);
            acemesh_task_set_type(STENCIL_TASK);
            acemesh_end_split_task();
        }
    }
    acemesh_spawn_and_wait(0);
    // acemesh_wait_for_all_task();
    return 0;
}

void Myinit()
{
    // int i,j,k;
    for(int h=0; h< MULTIG; h++)
    for (int i = 0; i < SIZEX; i++)
        for (int j = 0; j < SIZEY; j++)
            for (int k = 0; k < SIZEZ; k++)
            {
                double temp = (double)(i * 2.5 + j * 3.3 + k * 0.5) / 3.0;
                A[h][0][i][j][k] = temp;
                A[h][1][i][j][k] = temp;
                test[h][0][i][j][k] = temp;
                test[h][1][i][j][k] = temp;
            }
}
#ifdef INIT_CYCLIC
void Myinit_cyclic_dag(const int blocsizex, const int blocsizey)
{
    int i, j, k;
    for (int i = 0; i < SIZEX; i++)
        for (int j = 0; j < SIZEY; j++)
            for (int k = 0; k < SIZEZ; k++)
            {
                test[1][i][j][k] = (double)(i * 2.5 + j * 3.3 + k * 0.5) / 3.0;
                test[0][i][j][k] = test[1][i][j][k];
            }
    targs each;
    acemesh_begin_split_task("3d7p init_cyclic");
    for (i = 1; i < SIZEX - 1; i += blocsizex)
    {
        int starti = (i == 1) ? 0 : i;
        // int endi = (i + blocsizex > SIZEX -1 ) ? SIZEX - 1 : i + blocsizex;
        int endi = (i + blocsizex >= SIZEX - 1) ? SIZEX : i + blocsizex;
        for (j = 1; j < SIZEY - 1; j += blocsizey)
        {
            int startj = (j == 1) ? 0 : j;
            // int endj = (j + blocsizey > SIZEY -1 ) ? SIZEY - 1 : j + blocsizey;
            int endj = (j + blocsizey >= SIZEY - 1) ? SIZEY : j + blocsizey;

            each.starti = starti;
            each.endi = endi;
            each.startj = startj;
            each.endj = endj;

            acemesh_push_wlist(1, &A[0][i][j][0], NORMAL);
            acemesh_push_wlist(1, &A[1][i][j][0], NORMAL);
            acemesh_task_generator((TASK_FUNCPTR)stencil_init_cyclic,
                                   (void *)(&each), sizeof(targs));
            acemesh_task_set_type(STENCIL_TASK);
        }
    }
    acemesh_end_split_task();
    // acemesh_spawn_and_wait(1);
    /*    //init boundary data of array[0]
        for(int i = 0; i < SIZEX; i+=SIZEX-1)
        for(int j = 0; j < SIZEY; j++)//=SIZEY-1)
        for(int k = 0; k < SIZEZ; k++)//=SIZEZ-1)
        {
          A[0][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0;
          A[1][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0;
        }
        for(int i = 1; i < SIZEX-1; i++)//=SIZEX-1)
        for(int j = 0; j < SIZEY; j+=SIZEY-1)
        for(int k = 0; k < SIZEZ; k++)//=SIZEZ-1)
        {
          A[0][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0;
          A[1][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0;
        }
        for(int i = 1; i < SIZEX-1; i++)//=SIZEX-1)
        for(int j = 1; j < SIZEY-1; j++)//+=SIZEY-1)
        for(int k = 0; k < SIZEZ; k+=SIZEZ-1)
        {
          A[0][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0;
          A[1][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0;
        }*/
}
#endif
#ifdef INIT_BLOCKED
void Myinit_blocked_dag(int blocsizex, int blocsizey, int num_threads)
{
    int i, j, k;
    for (i = 0; i < SIZEX; i++)
        for (j = 0; j < SIZEY; j++)
            for (k = 0; k < SIZEZ; k++)
            {
                test[1][i][j][k] = (double)(i * 2.5 + j * 3.3 + k * 0.5) / 3.0;
                test[0][i][j][k] = test[1][i][j][k];
            }
    targs each;
    acemesh_arraytile(&A[1][0][0][0], "A", SIZEX, blocsizex, 2);
    acemesh_arraytile(&A[0][0][0][0], "A", SIZEX, blocsizex, 2);
    acemesh_begin_split_task("3d7p init_blocked");
    for (i = 1; i < SIZEX - 1; i += blocsizex)
    {
        int starti = (i == 1) ? 0 : i;
        // int endi = (i + blocsizex > SIZEX-1  ) ? SIZEX-1  : i + blocsizex;
        int endi = (i + blocsizex >= SIZEX - 1) ? SIZEX : i + blocsizex;
        int indexi = i / blocsizex;

        for (j = 1; j < SIZEY - 1; j += blocsizey)
        {
            acemesh_push_wlist(1, &A[0][i][j][0], NORMAL);
            acemesh_push_wlist(1, &A[1][i][j][0], NORMAL);
        }
        each.starti = starti;
        each.endi = endi;

        acemesh_task_generator((TASK_FUNCPTR)stencil_init_blocked,
                               (void *)(&each), sizeof(targs));
        acemesh_task_set_type(STENCIL_TASK);
        acemesh_affinity_from_arraytile(indexi, ceil((double)(SIZEX - 1 - 1) / blocsizex), num_threads);
    }
    acemesh_end_split_task();
    acemesh_spawn_and_wait(1);
    // acemesh_end_win();
    /*    //init boundary data of array[0]
        for(int i = 0; i < SIZEX; i+=SIZEX-1)
        for(int j = 0; j < SIZEY; j++)//=SIZEY-1)
        for(int k = 0; k < SIZEZ; k++)//=SIZEZ-1)
        {
          A[0][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0;
          A[1][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0;
        }
        for(int i = 1; i < SIZEX-1; i++)//=SIZEX-1)
        for(int j = 0; j < SIZEY; j+=SIZEY-1)
        for(int k = 0; k < SIZEZ; k++)//=SIZEZ-1)
        {
          A[0][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0;
          A[1][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0;
        }
        for(int i = 1; i < SIZEX-1; i++)//=SIZEX-1)
        for(int j = 1; j < SIZEY-1; j++)//+=SIZEY-1)
        for(int k = 0; k < SIZEZ; k+=SIZEZ-1)
        {
          A[0][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0;
          A[1][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0;
        }*/
}
#endif
int check(const int iter)
{
    int i, j, k, h;
    int itr;
    int ii, jj, kk;
    printf("begin check\n");

    for(h=0;h<MULTIG;h++)
    {
    for (itr = 0; itr < iter; ++itr)
    {
        for (i = 1; i < SIZEX - 1; i++)
            for (j = 1; j < SIZEY - 1; j++)
                for (k = 1; k < SIZEZ - 1; k++)
                    TestDo3d7pkernel(h,itr, i, j, k);
    }
    }

    for(h=0;h<MULTIG;h++)
    {
    for (ii = 1; ii < SIZEX - 1; ii++)
        for (jj = 1; jj < SIZEY - 1; jj++)
            for (kk = 1; kk < SIZEZ - 1; kk++)
            {
                if (fabs(test[h][(iter + 1) % 2][ii][jj][kk] - A[h][(iter + 1) % 2][ii][jj][kk]) > 1e-16)
                {
                    printf("h:%d, i: %d   j:  %d k: %d\n",h, ii, jj, kk);
                    printf("test: %lf", test[h][(iter + 1) % 2][ii][jj][kk]);
                    printf("but A: %lf", A[h][(iter + 1) % 2][ii][jj][kk]);
                    printf("they are different!\n");
                    return -1;
                }
            }
    }
    printf("correct\n");
    return 0;
}

int main(int argc, char **argv)
{
    int num_threads;
    // int i,j,k;
    struct timeval start;
    struct timeval end;
    double total_time;

    if (argc > 5)
    {
        printf("input arg error, USAGE: ./a.out nter blkx blky num_threads\n");
        return 0;
    }
    num_threads = 1;
    BLKX = 16;
    BLKY = 16;
    int ITER = 1;
    if (argc > 1)
        ITER = atoi(argv[1]);
    if (argc > 2)
        BLKX = atoi(argv[2]);
    if (argc > 3)
        BLKY = atoi(argv[3]);
    if (argc > 4)
        num_threads = atoi(argv[4]);
    acemesh_runtime_init(num_threads);
    // acemesh_sqrt_tasknum(256 / BLKX);
    // printf("set sqrt(tasknum) to:%d\n", 256 / BLKX);
    // acemesh_start_win();
    // gettimeofday(&start,NULL);
#ifdef INIT_SERIAL
    Myinit();
#elif INIT_CYCLIC
    Myinit_cyclic_dag(BLKX, BLKY);
#elif INIT_BLOCKED
    Myinit_blocked_dag(BLKX, BLKY, num_threads);
#endif

    gettimeofday(&start, NULL);
#ifdef INIT_BLOCKED
    Do3d7p(ITER, BLKX, BLKY, num_threads);
#else
    Do3d7p(ITER, BLKX, BLKY);
#endif
    gettimeofday(&end, NULL);

    acemesh_runtime_shutdown();
    total_time = 0.0;
    total_time += (end.tv_sec - start.tv_sec) + (end.tv_usec - start.tv_usec) / 1000000.0;
    printf("total time: %lf\n", total_time);
#ifdef CHECK
    if (check(ITER))
    {
        printf("check_error\n");
        exit(1);
    }
#endif
    return 0;
}
